/**
 * Copyright (C) 2024-2030 Super BilGen & 超亿时代信息技术股份有限公司.All Rights Reserved.
 */
package cn.com.yang.threadpool;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReentrantLock;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 
 * @Model cn.com.yang.threadpool.CustomThreadPool
 * @Desc TODO
 * @author yangyijin
 * @since 2025年9月4日-下午8:07:50
 */
public class CustomThreadPool {
	private static final Logger logger = LoggerFactory.getLogger(CustomThreadPool.class);

	// 线程池状态
	public enum PoolState {
		RUNNING, SHUTDOWN, TERMINATED
	}

	private volatile PoolState state = PoolState.RUNNING;

	// 核心线程数
	private final int corePoolSize;
	// 最大线程数
	private final int maxPoolSize;
	// 任务队列
	private final BlockingQueue<Runnable> taskQueue;
	// 工作线程集合
	private final List<WorkerThread> workers = new ArrayList<>();
	// 拒绝策略
	private final RejectionPolicy rejectionPolicy;
	// 线程工厂
	private final ThreadFactory threadFactory;
	// 空闲线程存活时间
	private final long keepAliveTime;
	// 工作线程计数器
	private final AtomicInteger threadCount = new AtomicInteger(0);
	// 提交任务计数器
	private final AtomicInteger submittedTaskCount = new AtomicInteger(0);
	// 完成任务计数器
	private final AtomicInteger completedTaskCount = new AtomicInteger(0);
	// 关闭标志
	private final AtomicBoolean isShutdown = new AtomicBoolean(false);
	// 全局锁，用于保护workers集合
	private final ReentrantLock mainLock = new ReentrantLock();

	// 构造函数
	public CustomThreadPool(int corePoolSize, int maxPoolSize, int queueCapacity, long keepAliveTime, TimeUnit unit,
			RejectionPolicy rejectionPolicy, ThreadFactory threadFactory) {
		this.corePoolSize = corePoolSize;
		this.maxPoolSize = maxPoolSize;
		this.taskQueue = new LinkedBlockingQueue<>(queueCapacity);
		this.keepAliveTime = unit.toNanos(keepAliveTime);
		this.rejectionPolicy = rejectionPolicy;
		this.threadFactory = threadFactory;

		logger.info("Initializing thread pool: core={}, max={}, queue={}, keepAlive={} {}", corePoolSize, maxPoolSize,
				queueCapacity, keepAliveTime, unit);

		// 初始化核心线程
		for (int i = 0; i < corePoolSize; i++) {
			addWorker(null);
		}

		logger.info("Thread pool initialized with {} core threads", corePoolSize);
	}

	// 简化构造函数 - 使用默认线程工厂和拒绝策略
	public CustomThreadPool(int corePoolSize, int maxPoolSize, int queueCapacity) {
		this(corePoolSize, maxPoolSize, queueCapacity, 60, TimeUnit.SECONDS, new AbortPolicy(),
				new DefaultThreadFactory("custom-pool"));
	}

	// 简化构造函数 - 可指定线程名前缀
	public CustomThreadPool(int corePoolSize, int maxPoolSize, int queueCapacity, String threadNamePrefix) {
		this(corePoolSize, maxPoolSize, queueCapacity, 60, TimeUnit.SECONDS, new AbortPolicy(),
				new DefaultThreadFactory(threadNamePrefix));
	}

	// 提交任务
	public void execute(Runnable task) {
		if (task == null)
			throw new NullPointerException();
		if (isShutdown.get()) {
			throw new IllegalStateException("ThreadPool is shutdown");
		}

		submittedTaskCount.incrementAndGet();
		logger.debug("Task submitted: {}", task);

		// 1. 如果工作线程数小于核心线程数，创建新线程
		if (threadCount.get() < corePoolSize) {
			if (addWorker(task)) {
				return;
			}
		}

		// 2. 尝试将任务加入队列
		if (taskQueue.offer(task)) {
			logger.debug("Task added to queue: {}", task);

			// 检查是否需要添加非核心线程处理积压任务
			if (threadCount.get() < maxPoolSize && taskQueue.size() > corePoolSize) {
				addWorker(null);
			}
			return;
		}

		// 3. 队列已满，尝试创建非核心线程
		if (threadCount.get() < maxPoolSize) {
			if (addWorker(task)) {
				return;
			}
		}

		// 4. 所有策略失败，执行拒绝策略
		logger.warn("Thread pool is saturated, executing rejection policy for task: {}", task);
		rejectionPolicy.reject(task, this);
	}

	// 添加工作线程
	private boolean addWorker(Runnable firstTask) {
		if (threadCount.get() >= maxPoolSize) {
			return false;
		}

		WorkerThread worker = new WorkerThread(firstTask);
		Thread thread = threadFactory.newThread(worker);
		if (thread != null) {
			worker.thread = thread;

			// 使用锁保护workers集合
			mainLock.lock();
			try {
				workers.add(worker);
			} finally {
				mainLock.unlock();
			}

			threadCount.incrementAndGet();
			thread.start();

			logger.debug("Worker thread created: {} for task: {}", thread.getName(), firstTask);
			return true;
		}
		return false;
	}

	// 关闭线程池（等待任务完成）
	public void shutdown() {
		if (isShutdown.compareAndSet(false, true)) {
			state = PoolState.SHUTDOWN;
			logger.info("Shutdown initiated. Waiting for tasks to complete...");

			// 中断空闲线程
			mainLock.lock();
			try {
				for (WorkerThread worker : workers) {
					if (!worker.isWorking.get()) {
						logger.trace("Interrupting idle thread: {}", worker.thread.getName());
						worker.thread.interrupt();
					}
				}
			} finally {
				mainLock.unlock();
			}
		}
	}

	// 立即关闭线程池
	public List<Runnable> shutdownNow() {
		if (isShutdown.compareAndSet(false, true)) {
			state = PoolState.SHUTDOWN;
			logger.warn("ShutdownNow initiated. Interrupting all threads...");

			List<Runnable> remainingTasks = new ArrayList<>();
			int drained = taskQueue.drainTo(remainingTasks);
			logger.info("Drained {} tasks from queue", drained);

			// 中断所有线程
			mainLock.lock();
			try {
				for (WorkerThread worker : workers) {
					logger.trace("Interrupting thread: {}", worker.thread.getName());
					worker.thread.interrupt();
				}
			} finally {
				mainLock.unlock();
			}
			return remainingTasks;
		}
		return new ArrayList<>();
	}

	// 等待线程池终止
	public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
		long nanos = unit.toNanos(timeout);
		long deadline = System.nanoTime() + nanos;

		synchronized (this) {
			while (state != PoolState.TERMINATED) {
				if (nanos <= 0L) {
					logger.warn("Timeout while waiting for termination");
					return false;
				}
				long remaining = deadline - System.nanoTime();
				wait(TimeUnit.NANOSECONDS.toMillis(remaining));
			}
			logger.info("Thread pool terminated successfully");
			return true;
		}
	}

	// 工作线程内部类
	private class WorkerThread implements Runnable {
		final AtomicBoolean isWorking = new AtomicBoolean(false);
		Runnable firstTask;
		Thread thread;
		final Logger workerLogger = LoggerFactory.getLogger(WorkerThread.class);

		WorkerThread(Runnable firstTask) {
			this.firstTask = firstTask;
		}

		@Override
		public void run() {
			Runnable task = firstTask;
			firstTask = null;

			try {
				workerLogger.debug("Worker thread started: {}", Thread.currentThread().getName());

				while (!Thread.currentThread().isInterrupted()) {
					// 执行当前任务
					if (task != null) {
						runTask(task);
						task = null;
					}

					// 从队列获取任务
					try {
						// 所有线程都使用超时等待，避免永久阻塞
						task = taskQueue.poll(keepAliveTime, TimeUnit.NANOSECONDS);

						// 如果超时且线程数大于核心数，回收线程
						if (task == null && threadCount.get() > corePoolSize) {
							workerLogger.debug("Thread idle timeout, exiting");
							break;
						}
					} catch (InterruptedException e) {
						workerLogger.debug("Thread interrupted while waiting for task");
						break;
					}
				}
			} finally {
				// 从工作线程集合中移除
				mainLock.lock();
				try {
					workers.remove(this);
				} finally {
					mainLock.unlock();
				}

				threadCount.decrementAndGet();
				workerLogger.debug("Worker thread exiting: {}", Thread.currentThread().getName());

				// 检查是否所有线程都已终止
				if (threadCount.get() == 0 && isShutdown.get()) {
					state = PoolState.TERMINATED;
					synchronized (CustomThreadPool.this) {
						logger.info("All worker threads terminated");
						CustomThreadPool.this.notifyAll();
					}
				}
			}
		}

		private void runTask(Runnable task) {
			isWorking.set(true);
			try {
				workerLogger.debug("Executing task: {}", task);
				task.run();
				workerLogger.debug("Task completed: {}", task);

				// 增加完成任务计数
				completedTaskCount.incrementAndGet();
			} catch (Exception e) {
				workerLogger.error("Task execution failed: {}", task, e);
			} finally {
				isWorking.set(false);
			}
		}
	}

	// 获取线程池状态信息
	public String getStatus() {
		int activeThreads = 0;
		for (WorkerThread worker : workers) {
			if (worker.isWorking.get()) {
				activeThreads++;
			}
		}

		return String.format("Threads: %d/%d (active: %d), Tasks: submitted=%d, completed=%d, queued=%d, State: %s",
				threadCount.get(), maxPoolSize, activeThreads, submittedTaskCount.get(), completedTaskCount.get(),
				taskQueue.size(), state);
	}

	// 获取活动线程数
	public int getActiveThreadCount() {
		int count = 0;
		for (WorkerThread worker : workers) {
			if (worker.isWorking.get()) {
				count++;
			}
		}
		return count;
	}

	// 获取提交任务数
	public int getSubmittedTaskCount() {
		return submittedTaskCount.get();
	}

	// 获取完成任务数
	public int getCompletedTaskCount() {
		return completedTaskCount.get();
	}

	// 获取队列任务数
	public int getQueuedTaskCount() {
		return taskQueue.size();
	}

	// 获取线程池状态
	public PoolState getState() {
		return state;
	}
}